import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import TimeSeriesSplit
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import warnings
warnings.filterwarnings('ignore')

plt.rcParams['font.sans-serif'] = ['WenQuanYi Micro Hei', 'Noto Sans CJK JP', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False  

class LotteryDataPreprocessor:
    def __init__(self, csv_path):
        self.df = pd.read_csv(csv_path)
        self.preprocess_data()
    
    def preprocess_data(self):
        """数据预处理"""
        # 重命名列
        self.df.columns = ['period', 'date', 'red_balls', 'blue_ball', 
                          'first_prize_count', 'first_prize_amount', 
                          'second_prize_count', 'second_prize_amount',
                          'sales_amount', 'prize_pool']
        
        # 转换时间列
        self.df['date'] = pd.to_datetime(self.df['date'].str.extract(r'(\d{4}-\d{2}-\d{2})')[0])
        self.df = self.df.sort_values('date').reset_index(drop=True)
        
        # 解析红球字符串为列表
        self.df['red_balls_list'] = self.df['red_balls'].str.split().apply(lambda x: [int(num) for num in x])
        
        # 创建单独的红球列
        for i in range(6):
            self.df[f'red{i+1}'] = self.df['red_balls_list'].apply(lambda x: x[i])
        
        # 创建时间特征
        self.df['year'] = self.df['date'].dt.year
        self.df['month'] = self.df['date'].dt.month
        self.df['day'] = self.df['date'].dt.day
        self.df['day_of_week'] = self.df['date'].dt.dayofweek
        self.df['day_of_year'] = self.df['date'].dt.dayofyear
        self.df['is_weekend'] = self.df['day_of_week'].isin([5, 6]).astype(int)
        
        # 创建号码统计特征
        self.create_number_features()
        
        # 创建奖金相关特征
        self.create_prize_features()
        
        return self.df
    
    def create_number_features(self):
        """创建号码相关的统计特征"""
        red_cols = ['red1', 'red2', 'red3', 'red4', 'red5', 'red6']
        
        # 基本统计
        self.df['red_sum'] = self.df[red_cols].sum(axis=1)
        self.df['red_mean'] = self.df[red_cols].mean(axis=1)
        self.df['red_std'] = self.df[red_cols].std(axis=1)
        self.df['red_median'] = self.df[red_cols].median(axis=1)
        self.df['red_range'] = self.df[red_cols].max(axis=1) - self.df[red_cols].min(axis=1)
        
        # 奇偶比
        self.df['red_odd_count'] = (self.df[red_cols] % 2 == 1).sum(axis=1)
        self.df['red_even_count'] = 6 - self.df['red_odd_count']
        self.df['red_odd_ratio'] = self.df['red_odd_count'] / 6
        
        # 大小比（以17为界）
        self.df['red_big_count'] = (self.df[red_cols] > 17).sum(axis=1)
        self.df['red_small_count'] = 6 - self.df['red_big_count']
        self.df['red_big_ratio'] = self.df['red_big_count'] / 6
        
        # 质数数量
        def is_prime(n):
            if n < 2:
                return False
            for i in range(2, int(np.sqrt(n)) + 1):
                if n % i == 0:
                    return False
            return True
        
        prime_count = []
        for _, row in self.df.iterrows():
            count = sum(is_prime(x) for x in row[red_cols])
            prime_count.append(count)
        self.df['red_prime_count'] = prime_count
        
        # 区间分布（1-11, 12-22, 23-33）
        self.df['red_low_count'] = (self.df[red_cols] <= 11).sum(axis=1)
        self.df['red_mid_count'] = ((self.df[red_cols] > 11) & (self.df[red_cols] <= 22)).sum(axis=1)
        self.df['red_high_count'] = (self.df[red_cols] > 22).sum(axis=1)
        
        # 连号数量
        def count_consecutive(numbers):
            sorted_nums = sorted(numbers)
            consecutive_count = 0
            for i in range(len(sorted_nums)-1):
                if sorted_nums[i+1] - sorted_nums[i] == 1:
                    consecutive_count += 1
            return consecutive_count
        
        self.df['consecutive_count'] = self.df[red_cols].apply(lambda row: count_consecutive(row.values), axis=1)
        
        # 创建滞后特征
        self.create_lag_features()
    
    def create_prize_features(self):
        """创建奖金相关特征"""
        # 奖金金额转换（提取数字部分并去除逗号）
        self.df['first_prize_amount'] = self.df['first_prize_amount'].astype(str).str.extract(r'([\d,]+)')[0].str.replace(',', '').astype(float)
        self.df['second_prize_amount'] = self.df['second_prize_amount'].astype(str).str.extract(r'([\d,]+)')[0].str.replace(',', '').astype(float)
        self.df['sales_amount'] = self.df['sales_amount'].astype(str).str.extract(r'([\d,]+)')[0].str.replace(',', '').astype(float)
        self.df['prize_pool'] = self.df['prize_pool'].astype(str).str.extract(r'([\d,]+)')[0].str.replace(',', '').astype(float)
        
        # 创建新特征
        self.df['prize_per_winner_first'] = self.df['first_prize_amount'] / self.df['first_prize_count'].replace(0, 1)
        self.df['prize_per_winner_second'] = self.df['second_prize_amount'] / self.df['second_prize_count'].replace(0, 1)
        
        self.df['sales_growth'] = self.df['sales_amount'].pct_change()
        self.df['prize_pool_growth'] = self.df['prize_pool'].pct_change()
        
        # 中奖难度指标
        self.df['winning_difficulty_first'] = self.df['first_prize_count'] / (self.df['sales_amount'] / 2)  # 每注2元
        self.df['winning_difficulty_second'] = self.df['second_prize_count'] / (self.df['sales_amount'] / 2)
    
    def create_lag_features(self, lags=5):
        """创建滞后特征"""
        features_to_lag = [
            'red_sum', 'red_mean', 'red_std', 'red_odd_count', 'red_big_count',
            'blue_ball', 'first_prize_count', 'sales_amount', 'prize_pool'
        ]
        
        for feature in features_to_lag:
            for lag in range(1, lags + 1):
                self.df[f'{feature}_lag_{lag}'] = self.df[feature].shift(lag)
        
        # 删除包含NaN的行
        self.df = self.df.dropna()
        
        return self.df
    
    def get_feature_matrix(self, target_type='numbers'):
        """获取特征矩阵"""
        if target_type == 'numbers':
            feature_cols = [
                'month', 'day', 'day_of_week', 'is_weekend',
                'red_sum', 'red_mean', 'red_std', 'red_range',
                'red_odd_count', 'red_big_count', 'red_prime_count',
                'red_low_count', 'red_mid_count', 'red_high_count',
                'consecutive_count', 'blue_ball',
                'sales_amount', 'prize_pool'
            ]
        elif target_type == 'prize':
            feature_cols = [
                'month', 'day', 'day_of_week', 'is_weekend',
                'red_sum', 'red_mean', 'red_std', 'red_range',
                'red_odd_count', 'red_big_count', 'consecutive_count',
                'blue_ball', 'sales_amount', 'prize_pool',
                'sales_growth', 'prize_pool_growth'
            ]
        
        # 添加所有滞后特征
        lag_features = [col for col in self.df.columns if '_lag_' in col]
        feature_cols.extend(lag_features)
        
        return self.df[feature_cols]
    
    def get_targets(self, target_type='numbers'):
        """获取目标变量"""
        if target_type == 'numbers':
            # 预测下一期的号码
            red_cols = ['red1', 'red2', 'red3', 'red4', 'red5', 'red6']
            targets = {}
            for col in red_cols:
                targets[col] = self.df[col].shift(-1)
            targets['blue'] = self.df['blue_ball'].shift(-1)
            
        elif target_type == 'prize':
            # 预测下一期的奖金相关信息
            targets = {
                'first_prize_count_next': self.df['first_prize_count'].shift(-1),
                'prize_per_winner_first_next': self.df['prize_per_winner_first'].shift(-1),
                'sales_amount_next': self.df['sales_amount'].shift(-1)
            }
        
        # 删除最后一行（没有下一期数据）
        target_df = pd.DataFrame(targets).iloc[:-1]
        
        return target_df

class LotteryPredictor:
    def __init__(self, preprocessor):
        self.preprocessor = preprocessor
        self.X = preprocessor.get_feature_matrix().iloc[:-1]  # 去掉最后一行
        self.y_red = preprocessor.get_targets()[['red1', 'red2', 'red3', 'red4', 'red5', 'red6']]
        self.y_blue = preprocessor.get_targets()['blue']
        
        # 数据标准化
        self.scaler_X = StandardScaler()
        self.scaler_y_red = StandardScaler()
        self.scaler_y_blue = StandardScaler()
        
        self.X_scaled = self.scaler_X.fit_transform(self.X)
        self.y_red_scaled = self.scaler_y_red.fit_transform(self.y_red)
        self.y_blue_scaled = self.scaler_y_blue.fit_transform(self.y_blue.values.reshape(-1, 1))
    
    def create_lstm_model(self, input_shape, output_units):
        """创建LSTM模型"""
        model = keras.Sequential([
            layers.LSTM(128, return_sequences=True, input_shape=input_shape),
            layers.Dropout(0.3),
            layers.LSTM(64, return_sequences=False),
            layers.Dropout(0.3),
            layers.Dense(32, activation='relu'),
            layers.Dense(output_units, activation='linear')
        ])
        
        model.compile(
            optimizer='adam',
            loss='mse',
            metrics=['mae']
        )
        
        return model
    
    def train_models(self, seq_length=10):
        """训练预测模型"""
        # 准备序列数据
        X_seq, y_red_seq, y_blue_seq = self.prepare_sequences(seq_length)
        
        # 划分训练测试集
        split_idx = int(len(X_seq) * 0.8)
        
        X_train, X_test = X_seq[:split_idx], X_seq[split_idx:]
        y_red_train, y_red_test = y_red_seq[:split_idx], y_red_seq[split_idx:]
        y_blue_train, y_blue_test = y_blue_seq[:split_idx], y_blue_seq[split_idx:]
        
        # 训练红球模型
        print("训练红球预测模型...")
        self.model_red = self.create_lstm_model(
            (seq_length, X_train.shape[2]), 
            6  # 预测6个红球
        )
        
        history_red = self.model_red.fit(
            X_train, y_red_train,
            epochs=100,
            batch_size=32,
            validation_data=(X_test, y_red_test),
            verbose=1,
            callbacks=[
                keras.callbacks.EarlyStopping(patience=10, restore_best_weights=True),
                keras.callbacks.ReduceLROnPlateau(factor=0.5, patience=5)
            ]
        )
        
        # 训练蓝球模型
        print("训练蓝球预测模型...")
        self.model_blue = self.create_lstm_model(
            (seq_length, X_train.shape[2]), 
            1  # 预测1个蓝球
        )
        
        history_blue = self.model_blue.fit(
            X_train, y_blue_train,
            epochs=100,
            batch_size=32,
            validation_data=(X_test, y_blue_test),
            verbose=1,
            callbacks=[
                keras.callbacks.EarlyStopping(patience=10, restore_best_weights=True),
                keras.callbacks.ReduceLROnPlateau(factor=0.5, patience=5)
            ]
        )
        
        return history_red, history_blue
    
    def prepare_sequences(self, seq_length):
        """准备时间序列数据"""
        X_seq, y_red_seq, y_blue_seq = [], [], []
        
        for i in range(len(self.X_scaled) - seq_length):
            X_seq.append(self.X_scaled[i:i+seq_length])
            y_red_seq.append(self.y_red_scaled[i+seq_length])
            y_blue_seq.append(self.y_blue_scaled[i+seq_length])
        
        return np.array(X_seq), np.array(y_red_seq), np.array(y_blue_seq)
    
    def predict_next(self, seq_length=10):
        """预测下一期号码"""
        # 获取最近seq_length期的数据
        recent_data = self.X_scaled[-seq_length:]
        recent_data = recent_data.reshape(1, seq_length, -1)
        
        # 预测
        pred_red_scaled = self.model_red.predict(recent_data)
        pred_blue_scaled = self.model_blue.predict(recent_data)
        
        # 反标准化
        pred_red = self.scaler_y_red.inverse_transform(pred_red_scaled)
        pred_blue = self.scaler_y_blue.inverse_transform(pred_blue_scaled)
        
        # 处理预测结果（号码应该是整数且在合理范围内）
        pred_red_processed = np.round(np.clip(pred_red[0], 1, 33)).astype(int)
        pred_blue_processed = np.round(np.clip(pred_blue[0][0], 1, 16)).astype(int)
        
        return sorted(pred_red_processed), pred_blue_processed

class LotteryAnalyzer:
    def __init__(self, df):
        self.df = df
    
    def analyze_basic_statistics(self):
        """基础统计分析"""
        print("=" * 60)
        print("双色球全面统计分析")
        print("=" * 60)
        
        red_cols = ['red1', 'red2', 'red3', 'red4', 'red5', 'red6']
        
        # 所有号码的出现频率
        all_red_numbers = pd.concat([self.df[col] for col in red_cols])
        blue_numbers = self.df['blue_ball']
        
        print(f"\n1. 红球出现频率统计：")
        red_freq = all_red_numbers.value_counts().sort_values(ascending=False)
        print(red_freq.head(15))
        
        print(f"\n2. 蓝球出现频率统计：")
        blue_freq = blue_numbers.value_counts().sort_values(ascending=False)
        print(blue_freq)
        
        print(f"\n3. 奖金相关统计：")
        print(f"平均一等奖注数: {self.df['first_prize_count'].mean():.2f}")
        print(f"平均一等奖金额: {self.df['prize_per_winner_first'].mean():.2f}元")
        print(f"平均销售额: {self.df['sales_amount'].mean():.2f}元")
        print(f"平均奖池金额: {self.df['prize_pool'].mean():.2f}元")
        
        return red_freq, blue_freq
    
    def plot_comprehensive_analysis(self):
        """综合可视化分析"""
        fig, axes = plt.subplots(4, 2, figsize=(16, 16))
        
        # 号码趋势
        self.df['red_sum'].plot(ax=axes[0, 0], title='红球和随时间变化', color='red')
        axes[0, 0].set_ylabel('红球和')
        
        self.df['blue_ball'].plot(ax=axes[0, 1], title='蓝球号码趋势', color='blue')
        axes[0, 1].set_ylabel('蓝球号码')
        
        # 销售额和奖池
        self.df['sales_amount'].plot(ax=axes[1, 0], title='销售额趋势', color='green')
        axes[1, 0].set_ylabel('销售额(元)')
        
        self.df['prize_pool'].plot(ax=axes[1, 1], title='奖池金额趋势', color='orange')
        axes[1, 1].set_ylabel('奖池金额(元)')
        
        # 中奖情况
        self.df['first_prize_count'].plot(ax=axes[2, 0], title='一等奖注数趋势', color='purple')
        axes[2, 0].set_ylabel('注数')
        
        self.df['prize_per_winner_first'].plot(ax=axes[2, 1], title='单注一等奖金额趋势', color='brown')
        axes[2, 1].set_ylabel('金额(元)')
        
        # 号码特征
        self.df['red_odd_count'].plot(ax=axes[3, 0], title='红球奇数个数趋势', color='darkred')
        axes[3, 0].set_ylabel('奇数个数')
        
        self.df['consecutive_count'].plot(ax=axes[3, 1], title='连号数量趋势', color='darkblue')
        axes[3, 1].set_ylabel('连号数量')
        
        plt.tight_layout()
        plt.show()
    
    def analyze_patterns(self):
        """模式分析"""
        print("\n4. 号码模式分析：")
        
        # 奇偶模式
        odd_even_pattern = self.df[['red_odd_count', 'red_even_count']].mean()
        print(f"平均奇偶比: {odd_even_pattern['red_odd_count']:.1f}:{odd_even_pattern['red_even_count']:.1f}")
        
        # 大小模式
        big_small_pattern = self.df[['red_big_count', 'red_small_count']].mean()
        print(f"平均大小比: {big_small_pattern['red_big_count']:.1f}:{big_small_pattern['red_small_count']:.1f}")
        
        # 区间分布
        print(f"号码区间分布 - 低(1-11): {self.df['red_low_count'].mean():.2f}, "
              f"中(12-22): {self.df['red_mid_count'].mean():.2f}, "
              f"高(23-33): {self.df['red_high_count'].mean():.2f}")
        
        # 连号分析
        print(f"平均连号数量: {self.df['consecutive_count'].mean():.2f}")
        
        # 质数分析
        print(f"平均质数个数: {self.df['red_prime_count'].mean():.2f}")
    
    def analyze_market_trends(self):
        """市场趋势分析"""
        print("\n5. 市场趋势分析：")
        
        # 销售额分析
        sales_stats = self.df['sales_amount'].describe()
        print(f"销售额统计: 均值={sales_stats['mean']:.0f}, 最大值={sales_stats['max']:.0f}, 最小值={sales_stats['min']:.0f}")
        
        # 奖池分析
        prize_pool_stats = self.df['prize_pool'].describe()
        print(f"奖池统计: 均值={prize_pool_stats['mean']:.0f}, 最大值={prize_pool_stats['max']:.0f}")
        
        # 中奖难度
        print(f"一等奖中奖概率: {self.df['winning_difficulty_first'].mean():.8f}")
        print(f"二等奖中奖概率: {self.df['winning_difficulty_second'].mean():.6f}")
    
    def correlation_analysis(self):
        """相关性分析"""
        numeric_cols = self.df.select_dtypes(include=[np.number]).columns
        # 选择关键特征进行相关性分析
        key_features = [
            'red_sum', 'red_odd_count', 'red_big_count', 'consecutive_count',
            'blue_ball', 'first_prize_count', 'prize_per_winner_first',
            'sales_amount', 'prize_pool'
        ]
        
        correlation_matrix = self.df[key_features].corr()
        
        plt.figure(figsize=(12, 10))
        sns.heatmap(correlation_matrix, annot=True, cmap='coolwarm', center=0, fmt='.2f')
        plt.title('关键特征相关性矩阵')
        plt.tight_layout()
        plt.show()
        
        return correlation_matrix
    
    def plot_number_frequency(self):
        """号码频率分布图"""
        red_cols = ['red1', 'red2', 'red3', 'red4', 'red5', 'red6']
        all_red_numbers = pd.concat([self.df[col] for col in red_cols])
        
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))
        
        # 红球频率
        red_freq = all_red_numbers.value_counts().sort_index()
        ax1.bar(red_freq.index, red_freq.values, color='red', alpha=0.7)
        ax1.set_title('红球号码出现频率')
        ax1.set_xlabel

def main():
    # 1. 数据预处理
    preprocessor = LotteryDataPreprocessor('/home/ubuntu/mypro/other/ml-pro/project/data/ssq_data_20250912_075157.csv')
    df_processed = preprocessor.df
    
    # 2. 数据分析
    analyzer = LotteryAnalyzer(df_processed)
    red_freq, blue_freq = analyzer.analyze_basic_statistics()
    analyzer.plot_comprehensive_analysis()
    analyzer.plot_number_frequency()
    analyzer.analyze_patterns()
    corr_matrix = analyzer.correlation_analysis()
    
    # 3. 模型训练（如果有足够的数据）
    if len(df_processed) > 50:  # 至少需要50期数据
        predictor = LotteryPredictor(preprocessor)
        history_red, history_blue = predictor.train_models(seq_length=10)
        
        # 4. 预测下一期
        next_red, next_blue = predictor.predict_next()
        print(f"\n预测下一期号码：")
        print(f"红球: {next_red}")
        print(f"蓝球: {next_blue}")
    
    # 5. 保存处理后的数据
    df_processed.to_csv('processed_lottery_data.csv', index=False, encoding='utf-8-sig')
    print("\n处理后的数据已保存到 processed_lottery_data.csv")

if __name__ == "__main__":
    main()